import os
import numpy as np
from typing_extensions import Iterable


class DatabaseBuilder:
    """
    导航数据库构建基类
    """

    def __init__(self, path, radius=1737.4e3, cache_path=None, *args, **kwargs):
        """
        Initialize the DatabaseBuilder.

        This constructor sets up the database builder, loading or building the catalog
        based on the provided parameters.

        Parameters:
        path (str): The path to the input data file.
        radius (float, optional): The radius of the celestial body in meters.
                                  Defaults to 1737.4e3 (Moon's radius).
        cache_path (str, optional): Path to cache the built catalog. If provided and
                                    the cache exists, it will be loaded instead of
                                    rebuilding the catalog. Defaults to None.
        *args: Variable length argument list passed to the build method.
        **kwargs: Arbitrary keyword arguments passed to the build method.

        Returns:
        None

        Note:
        The method either loads an existing catalog from the cache, builds a new one
        and caches it, or builds a new one without caching, depending on the provided
        parameters and the existence of a cache file.
        """
        self.path = path
        self.radius = radius
        self.cache_path = cache_path
        if cache_path:
            if not os.path.exists(cache_path):
                data = self.load()
                catalog_dict = self.build(data.T, *args, **kwargs)
                np.savez(cache_path, **catalog_dict)
            else:
                catalog_dict = np.load(cache_path, allow_pickle=True)
                catalog_dict = {k: catalog_dict[k] for k in catalog_dict.files}
        else:
            data = self.load()
            catalog_dict = self.build(data.T, *args, **kwargs)
        self.data = catalog_dict

    def __init__(self, path, radius=1737.4e3, cache_path=None, *args, **kwargs):
        self.path = path
        self.radius = radius
        self.cache_path = cache_path
        if cache_path:
            if not os.path.exists(cache_path):
                data = self.load()
                catalog_dict = self.build(data.T, *args, **kwargs)
                np.savez(cache_path, **catalog_dict)
            else:
                catalog_dict = np.load(cache_path, allow_pickle=True)
                catalog_dict = {k: catalog_dict[k] for k in catalog_dict.files}
        else:
            data = self.load()
            catalog_dict = self.build(data.T, *args, **kwargs)
        self.data = catalog_dict

    def __call__(self, *args, **kwargs):
        if self.cache_path:
            if not os.path.exists(self.cache_path):
                data = self.load()
                catalog_dict = self.build(data.T, *args, **kwargs)
                np.savez(self.cache_path, **catalog_dict)
            else:
                catalog_dict = np.load(self.cache_path, allow_pickle=True)
                catalog_dict = {k: catalog_dict[k] for k in catalog_dict.files}
        else:
            data = self.load()
            catalog_dict = self.build(data.T, *args, **kwargs)

        return catalog_dict

    def build(self, data: Iterable, max_field=None) -> dict:
        """构建数据库的核心方法"""
        raise NotImplementedError

    def load(self):
        """
        注意，加载的数据集应当是一个逗号分隔文件，其中每一行的格式为：
        lat, lon, radius, depth
        应当保证经纬度是以度为单位的，半径和深度是以米为单位的
        """
        data_list = []
        with open(self.path, "r") as f:
            # 取消第一行
            f.readline()
            for line in f:
                # 按lat, lon, radius, depth的顺序返回
                data = [float(x) for x in line.split(",")]
                if len(data) == 3:
                    data.append(0)
                assert len(data) == 4
                data_list.append(data)
        return np.array(data_list)

    def to_csv(self, output_dir):
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        np.savetxt(
            os.path.join(output_dir, "catalog.csv"),
            self.data["catalog"],
            delimiter=",",
            fmt="%.4f",
        )
        np.savetxt(
            os.path.join(output_dir, f"descriptor.csv"),
            self.data["descriptor"].T,
            delimiter=",",
            fmt="%.4f",
        )
        np.savetxt(
            os.path.join(output_dir, f"idx.csv"),
            self.data["idx"].T,
            delimiter=",",
            fmt="%d",
        )
        np.savetxt(
            os.path.join(output_dir, "temp.csv"),
            self.data["temp"],
            delimiter=",",
            fmt="%.4f",
        )
